#lang racket

(require "util/utils.scm")
(provide pass-temps)

(define (pass-temps prg)
  (let ((definitions (cdr (assoc 'definitions prg)))
        (procedures (cdr (assoc 'procedures prg))))
    `((definitions . ,(map* pass-temps-definitions definitions))
      (procedures . ,(map* pass-temps-procedures procedures)))))

(define (pass-temps-definitions definition)
  definition)

(define (pass-temps-procedures procedures)
  (let ((label (car procedures))
        (env (cadr procedures))
        (binds (caddr procedures))
        (exp (cadddr procedures)))
    `(,label ,env ,binds ,(count-temps-let exp) ,exp)))

(define (count-temps-let exp)
  (match exp
    (`(let* ,binds ,body)
     (let loop ((binds binds) (temps 0))
       (if (null? binds)
	   (+ temps (count-temps-exp body))
	   (loop (cdr binds) (+ temps (count-temps-bind (car binds)))))))
    (else (count-temps-exp exp))))

(define (count-temps-bind bind)
  (match bind
    (`(,var ,exp) (+ 1 (count-temps-exp exp)))
    (else (error 'count-temps-bind "not a valid bind expression in let*" bind))))

(define (count-temps-exp exp)
  (match exp
    (`(atomic ,x) 0)
    (`(var ,type ,x) 0)
    (`(closure ,x ,y) 0)
    (`(quote ,x) 0)
    (`(builtin ,nm . ,args) 0)
    (`(app . ,vars) 0)
    (`(known-app ,i . ,vars) 0)
    (`(if ,t ,c ,a) (max (count-temps-let c) (count-temps-let a)))
    (else (error 'count-temps-exp "not a valid expression inside the body of a let*" exp))))
